# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 11:53 
@Author  : ZhangShenao 
@File    : chat_memory_chain.py 
@Desc    : 聊天记忆链

把聊天历史的记忆功能单独封装成一个Chain
"""
from operator import itemgetter
from typing import Dict

from langchain.memory import ConversationBufferWindowMemory
from langchain.memory.chat_memory import BaseChatMemory
from langchain_community.chat_message_histories import FileChatMessageHistory
from langchain_core.language_models import BaseChatModel
from langchain_core.output_parsers import BaseTransformOutputParser
from langchain_core.prompts import BaseChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableConfig
from langchain_core.runnables.utils import Output
from langchain_core.tracers import Run
from typing_extensions import Any

from internal.service import VectorStoreService

CHAT_HISTORY_FILE_PATH = '../../storage/memory/chat_history.json'  # 聊天历史文件路径
HISTORY_KEY = 'history'  # 聊天历史Key
CONTEXT_KEY = 'context'  # 上下文信息Key
INPUT_KEY = 'input'  # 聊天输入Key
OUTPUT_KEY = 'output'  # 聊天输出Key
MEMORY_CONFIG_KEY = 'memory'  # 记忆配置Key
SAVE_CONVERSATION_ROUNDS = 5  # 保存历史对话的轮数


def invoke_chain_with_chat_memory(input: Dict[str, Any],
                                  prompt_template: BaseChatPromptTemplate,
                                  llm: BaseChatModel,
                                  parser: BaseTransformOutputParser[str],
                                  vector_store_service: VectorStoreService) -> Output:
    """
    将传入的Runnable组件编排成Chain,并在此基础上封装聊天记忆功能,返回最终调用结果
    :param input: 调用输入参数
    :param prompt_template: 提示词模板
    :param llm: LLM聊天模型
    :param parser: 输出解析器
    :return: Chain调用结果
    """

    # 使用本地文件保存聊天历史
    chat_history = FileChatMessageHistory(file_path=CHAT_HISTORY_FILE_PATH)

    # 创建聊天记忆组件,使用缓冲窗口记忆方式
    memory = ConversationBufferWindowMemory(
        input_key=INPUT_KEY,
        output_key=OUTPUT_KEY,
        memory_key=HISTORY_KEY,
        k=SAVE_CONVERSATION_ROUNDS,  # 保留最近5轮的聊天历史,即10条消息
        return_messages=True,  # 结果返回聊天消息列表,而不是字符串
        chat_memory=chat_history  # 设置MessageHistory组件,用于持久化历史聊天记录
    )

    # 构建Retriever Chain,执行向量相似度检索
    retriever = vector_store_service.as_retriever() | vector_store_service.join_document_page_contents

    # 构造Chain执行链,用于编排组件的执行流程
    chain = RunnablePassthrough.assign(
        context=itemgetter(INPUT_KEY) | retriever,
        history=RunnableLambda(_load_memory_variables_from_config) | itemgetter(HISTORY_KEY)
    ) | prompt_template | llm | parser

    # 封装chain,在运行配置中传入记忆信息,并且注册on_end监听回调,在回调函数中保存聊天历史
    memory_chain = (chain.with_config(configurable={MEMORY_CONFIG_KEY: memory})
                    .with_listeners(on_end=_save_chat_history))

    # 调用memory_chain,返回结果
    output = memory_chain.invoke(input)
    return output


def _load_memory_variables_from_config(input: Dict[str, Any], config: RunnableConfig) -> Dict[str, Any]:
    """
    从运行配置中,加载记忆变量
    :param input: 运行调用输入
    :param config: 运行配置
    :return: 记忆变量字典
    """

    # 获取运行时配置,从配置读取记忆信息
    conf = config.get('configurable', {})
    memory = conf.get(MEMORY_CONFIG_KEY, None)
    if memory is not None and isinstance(memory, BaseChatMemory):
        return memory.load_memory_variables(input)

    # 空记忆信息
    return {}


def _save_chat_history(run_obj: Run, config: RunnableConfig) -> None:
    """
    保存聊天历史
    :param run_obj: 运行时对象,包含了所有运行时的相关信息
    :param config: 运行时配置信息
    """
    # 获取运行时配置,从配置读取记忆信息
    conf = config.get('configurable', {})
    memory = conf.get(MEMORY_CONFIG_KEY, None)
    if memory is not None and isinstance(memory, BaseChatMemory):
        # 将当前聊天的输入输出保存到Memory中
        memory.save_context(run_obj.inputs, run_obj.outputs)
